#!/usr/bin/env python

"""
Populates a MySQL database with a combination of ARB's and NCBI's taxonomy
databases.

"""
import ssummo.CONFIG as CONFIG

import ftplib
import os
import re
import subprocess
import sys
import tarfile

try:
    import cPickle as pickle
except ImportError:
    import pickle

import MySQLdb

class ISqlDatabase(object):
    """Abstract class providing database connectivity hooks for sublcasses to
    implement."""

    def __init__(self):
        """Initialise database variables"""
        super(ISqlDatabase, self).__init__()
        self.user = CONFIG.user  #: MySQL user name
        self.host = CONFIG.host  #: MySQL server host name
        self.passwd = CONFIG.passwd  #: MySQL password
        self.db = CONFIG.db  #: MySQL database name

    def connect(self, user, host, passwd, db):
        """Connect to MySQL database, creating :data:`self.cnx` and
        :data:`self.cur`"""
        print("Connecting to mysql database: {0}".format(db))
        self.cnx = MySQLdb.connect(user=user, host=host, passwd=passwd, db=db)
        self.cur = self.cnx.cursor()

    def delete_tables(self):
        """Drop all tables from database. Use with care!!!"""
        print('Deleting tables: {0}'.format(','.join(self.tables)))
        for table in self.tables:
            self.cur.execute('DROP TABLE {0};'.format(table))

    def show_tables(self):
        """Populate :data:`.tables` from the SQL command::

            (mysql) SHOW TABLES
        """
        self.cur.execute('SHOW TABLES;')
        rows = self.cur.fetchall()
        self.tables = []
        for row in rows:
            self.tables.append(row[0])


class NcbiTaxDb(ISqlDatabase):
    """Class used to retrieve the latest EMBL taxonomy listing
    from the NCBI ftp repository and save it to a local mysql
    database.

    _write_to_mysql may need to be edited to match the database
    security info that is in use on your system.

    General useage:
    x = NCBI_tax_to_SQL()
    x.automate()

    By default this will use the database metagenomics (will crash
    if the database doesn't exist) and then create the table
    'nodes' within that database.
    """
    def __init__(self):
        super(NcbiTaxDb, self).__init__()
        self.to_get = 'pub/taxonomy/taxdump.tar.gz'
        self.out_name = 'taxdump.tar.gz'
        self.connect(self.user, self.host, self.passwd, self.db)

    def _FTP_connect(self):
        print("logging in to NCBI's ftp servers...")
        self.ftp = ftplib.FTP('ftp.ncbi.nih.gov')
        self.ftp.login('anonymous', '')
#       return self

    def _get_tax_tarfile(self):
        print("Retrieving the compressed taxonomy database")
        with file(self.out_name, 'wb') as out_handle:
            self.ftp.retrbinary("RETR " + self.to_get, out_handle.write)
        self.ftp.quit()

    def _extract_nodes(self):
        print("Extracting and writing to 'nodes.dmp'")
        tarfile.open(self.out_name, 'r:gz').extract('nodes.dmp', path='.')
        print("Extracting and writing to 'names.dmp'")
        tarfile.open(self.out_name, 'r:gz').extract('names.dmp', path='.')

    def iter_lines(self, fname):
        with file(fname, 'r') as handle:
            for line in handle:
                yield line

    def _write_nodes_to_mysql(self, db_name='taxonomies'):
        print("Putting nodes.dmp into mysql database")
        self.show_tables()
        if 'nodes' in self.tables:
            self.cur.execute("DROP TABLE nodes;")
        self.cur.execute("CREATE TABLE nodes (tax_id INT NOT NULL PRIMARY KEY,\
                     parent_tax_id INT,\
                     rank VARCHAR(100),\
                     embl_code VARCHAR(100),\
                     FOREIGN KEY (parent_tax_id) references names(tax_ID),\
                     FOREIGN KEY (tax_ID) references names(tax_ID));")  ####,\
#                    division_ID INTEGER,\
#                    inherited_div_flag ENUM('0','1'),\
#                    gen_code_ID INTEGER,\
#                    inherited_GCFlag ENUM('0','1'),\
#                    mitochondrial_gen_code_ID INTEGER,\
#                    inherited_MGC_flag ENUM('0','1'),\
#                    Genbank_hidden_flag ENUM('0','1'),\
#                    hidden_subtree_root_flag ENUM('0','1'),\
#                    comments VARCHAR(200));")
        for line in self.iter_lines('nodes.dmp'):
            row = line.split('|')
            self.cur.execute("INSERT INTO nodes SET tax_ID='{0}',\
                     parent_tax_id='{1}',\
                     rank='{2}',\
                     embl_code='{3};'"\
                     .format(row[0].strip(),
                             row[1].strip(),
                             row[2].strip(),
                             row[3].strip()))  ###,\
#                    division_ID='{4}',\
#                    inherited_div_flag='{5}',\
#                    gen_code_ID='{6}',\
#                    inherited_GCFlag='{7}',\
#                    mitochondrial_gen_code_ID='{8}',\
#                    inherited_MGC_flag='{9}',\
#                    Genbank_hidden_flag='{10}',\
#                    hidden_subtree_root_flag='{11}',\
#                    comments='{12}';".format(row[0],row[1],row[2],row[3],row[4],row[5],row[6],row[7],row[8],row[9],row[10],row[11],row[12]))
        self.cnx.commit()

    def _write_names_to_mysql(self, db_name='taxonomies'):
        print("Putting names.dmp into mysql database")
        self.show_tables()
        if 'names' in self.tables:
            self.cur.execute("DROP TABLE names;")
        self.cur.execute("CREATE TABLE names (tax_ID int NOT NULL, name VARCHAR(200) NOT NULL, UniqueName VARCHAR(200), NameClass VARCHAR(50));")
        for line in self.iter_lines('names.dmp'):
            row = re.split('\|',
                           re.sub(r"""(['"])|((?<=\|)\s+|\s+(?=\|))|(\/)""",
                                  lambda m: '.' if m.groups()[2] else '', line))
            if not re.search('scientific', row[3], re.I):
                continue
            cmd = 'INSERT INTO names SET tax_ID="{0}",' + \
                  'name="{1}",UniqueName="{2}",NameClass="{3}";'\
                  .format(row[0], row[1], row[2], row[3])
            try:
                self.cur.execute(cmd)
            except Exception:
                print(cmd)
                raise
        self.cnx.commit()

    def _clean_up(self):
        os.remove('nodes.dmp')
        os.remove('names.dmp')
#       os.remove('taxdump.tar.gz')
        self.cnx.close()

    def automate(self, db_name='taxonomies'):
        if not os.path.exists(self.out_name) and not \
                (os.path.exists('nodes.dmp') and os.path.exists('names.dmp')):
            self._FTP_connect()
            self._get_tax_tarfile()
        self._extract_nodes()
        self._write_nodes_to_mysql()
        self._write_names_to_mysql()
        self._clean_up()
        print("finished successfully! =D")


class Translate(ISqlDatabase):
    def __init__(self):
        super(NcbiTaxDb, self).__init__()

        with file(os.path.join(CONFIG.top, CONFIG.taxIndex), 'rb') as in_handle:
            self.arb_dict = pickle.load(in_handle)

        self.connect()
        self.show_tables()
        self.create_output_table()

    def create_output_table(self):
        """Creates a MySQL table called EMBL_Vs_ARB and a text file
        called EMBL_Vs_ARB.txt.
        These are empty but can be written to by calling:-
        self.cur.execute( ... ) or
        self.out_handle.write( ... )
        Columns are in the order:
        OTU Name*, EMBL Count, ARB Count, Rank, Depth*, Max Depth
        The *'s indicate compound primary keys."""
        print("Initiating database table and text output file")
#       if 'EMBL_Vs_ARB' in self.tables or 'embl_vs_arb' in self.tables: ### Need lower case for windows mysql servers
#           self.cur.execute('DROP TABLE EMBL_Vs_ARB;')
#           self.show_tables()
        if 'Eukaryotes' in self.tables or 'eukaryotes' in self.tables: ### Need lower case for windows mysql servers
            self.cur.execute('DROP TABLE Eukaryotes;')
            self.show_tables()
        if 'Prokaryotes' in self.tables or 'prokaryotes' in self.tables: ### Need lower case for windows mysql servers
            self.cur.execute('DROP TABLE Prokaryotes;')
            self.show_tables()
        if 'NonUniques' in self.tables or 'nonuniques' in self.tables: ### Need lower case for windows mysql servers
            self.cur.execute('DROP TABLE NonUniques;')
            self.show_tables()
        self.cur.execute('CREATE TABLE Eukaryotes (Name VARCHAR(100), ParentName VARCHAR(100), tax_ID INT, Rank VARCHAR(50), PRIMARY KEY (Name, ParentName));')
        self.cur.execute('CREATE TABLE Prokaryotes (Name VARCHAR(100), ParentName VARCHAR(100), tax_ID INT, Rank VARCHAR(50), PRIMARY KEY (Name, ParentName));')
        self.cur.execute('CREATE TABLE NonUniques (Root CHAR(12),Name VARCHAR(100),ParentName VARCHAR(100), tax_IDs VARCHAR(100), PRIMARY KEY (Root,Name,ParentName));')

    def walk(self, dict_to_walk, parent='root'):
        """Provide a dictionary, and this will traverse it
        top down, yielding each OTU, the depth it is found
        and the maximum depth from that point."""
        node = dict_to_walk
        for OTU in node.keys():
            if OTU == 'accessions':
                continue
            yield OTU, parent
            for x in self.walk(node[OTU], parent=OTU):
                yield x

    def rows_to_list(self, rows):
        """Given a single column of results from mysql.cur.fetchall()
        convert to a single list"""
        row = []
        for r in rows:
            assert len(r) == 1
            row.append(r[0])
#       row = [r[0] for r in rows]
        return row

    def rows_to_set(self, rows):
        row = set()
        for r in rows:
            assert len(r) == 1
            row.add(r[0])
        return row

    def _multiple_parents(self, childIDs, parent, child):
        possible_parents = []
        self.cur.execute('SELECT tax_ID FROM names WHERE name="{0}" AND '+\
                         'NameClass="scientific name";'.format(parent))
        parent_IDs = self.rows_to_set(self.cur.fetchall())
        if len(parent_IDs) == 1:    ## If there is only 1 TaxID with the same parent name....
            parent_ID = parent_IDs.pop()
            ### Check that a node exists with this parent_tax_ID and one of the childIDs.
            for childID in childIDs:
                self.cur.execute('SELECT rank FROM nodes WHERE tax_ID={0} AND '+\
                                 'parent_tax_ID={1};'.format(childID, parent_ID))
                ranks = self.rows_to_set(self.cur.fetchall())
                if len(ranks) == 0:
                    continue
                else:
                    assert len(ranks) == 1
                    rank = ranks.pop()
                    break
            try:
                return childID, parent_ID, rank
            except UnboundLocalError:
                rank = 'unknown'
                return childID, parent_ID, rank
        elif len(parent_IDs)>1:
            parentKids = {}
            for parent_ID in parent_IDs:
                cmd = 'SELECT tax_ID FROM nodes WHERE parent_tax_ID={0}'\
                      .format(parent_ID)
                self.cur.execute(cmd)
                possKids = self.rows_to_set(self.cur.fetchall())
                for kid in possKids:
                    if kid in childIDs:
                        parentKids.update( { parent_ID : kid } )
            if len(parentKids) == 1:
                parentID, childID = parentKids.items()[0]
                rank = self.get_rank(childID)
                return childID, parentID, rank
            else:
                pass
        else:
            pass
        print("Couldn't find a match for the {0} IDs provided.".format(len(childIDs)))
        #print("Names: {0}.".format([n[0] for n in names]))
        print("possible parents: {0}".format(possible_parents))
        print("Parent name: {0}".format(parent))
        return (0, 0, 'unknown')

    def find_parent(self, childIDs, parent, child):
        """Given a list of ID numbers and the name of the parent
        this will find the ID with the same parent (from IDs) and
        will return the correct ID and the rank. """
        if len(childIDs) == 1:
            ID = childIDs.pop()
            self.cur.execute('SELECT parent_tax_id,rank FROM nodes where tax_ID="{0}";'.format(ID) )
            rows = self.cur.fetchall()
            if len(rows) == 1:
                parentID, rank = rows[0]
                self.cur.execute('SELECT name from names WHERE tax_ID="{0}" and NameClass="scientific name";'.format( parentID ))
                names = self.cur.fetchall()
                names = self.rows_to_list(names)
                if parent in names:
                    return parentID, rank
                try:
                    assert len(names) == 1
                except AssertionError:
                    print('Multiple names for parent with ID {0}'.format(parentID))
                    print(names+'\n')
                    raise IOError('Multiple entries for TaxID {0}'.format(parentID))
                return ID, parentID, rank
            elif len(rows) > 1:
                print("found multiple parent IDs for '{0}':-".format(child))
                print("{0}\nUsed ID: '{1}'".format(rows, ID))
                return self._multiple_parents(childIDs, parent, child)
            else:
                print("No matching parents for child '{0}' with taxID '{1}'"
                      .format(child, ID))
                return 0, 0, 'unknown'
        elif len(childIDs) > 1:
            return self._multiple_parents(childIDs, parent, child)

    def get_rank(self, ID):
        self.cur.execute('SELECT rank FROM nodes WHERE tax_ID={0}'.format(ID))
        rank = self.rows_to_set( self.cur.fetchall() )
        assert len(rank) == 1
        return rank.pop()

    def get_Ids(self, OTU, parent):
        self.cur.execute('SELECT tax_ID FROM names WHERE name="{0}";'.format(OTU))
        tax_IDs = self.cur.fetchall()
        tax_IDs = self.rows_to_set( tax_IDs)
        nIDs = len(tax_IDs)
        if nIDs == 1:
            childID = tax_IDs.pop()
            rank = self.get_rank(childID)
        elif nIDs > 1:
            print("Got {0} IDs".format(nIDs))
            try:
                childID, parentID, rank = self.find_parent(tax_IDs, parent, OTU)
            except Exception:
                print('table_IDs: {0}'.format( tax_IDs))
                print('parent: {0}\nOTU: {1}'.format(parent, OTU))
                raise
        else:
            childID = 0
            rank = "unknown"
        return str(int(childID)), rank

    def _get_non_unique_tax_Ids(self, root, OTU, parent):
        self.cur.execute('SELECT tax_IDs FROM NonUniques WHERE Root="{0}" AND '+
                         'Name="{1}" AND ParentName="{2}";'
                         .format(root, OTU, parent))
        result = self.cur.fetchall()
        exIDs = self.rows_to_list(result)
        ID, rank = self.get_Ids(OTU, parent)
        if exIDs == []:
            exIDs = [str(ID)]
            new = True
        else:

            exIDs = ' '.split(exIDs[0])
            if ID not in exIDs:
                exIDs.append(ID)
            new = False
        return exIDs, new

    def _update_all(self):
        print("Traversing the ARB database and updating the MySQLdb")
        print('\t|\t'.join(['Table', 'OTU Name', 'parent', 'tax_ID', 'Rank']))
        for root in ['Archaea', 'Bacteria', 'Eukaryota']:
            if root == 'Archaea' or root == 'Bacteria':
                table = 'Prokaryotes'
            else:
                table = 'Eukaryotes'
            if table in self.tables:
                self.cnx.execute('DROP TABLE {0};'.format(table))
            try:
                startNode = self.arb_dict[root]
            except KeyError:
                print(self.arb_dict.keys())
                raise
            ID, rank = self.get_Ids(root, 'root')
            self.cur.execute('INSERT INTO {0} SET Name="{1}",ParentName="{2}",'+
                             'tax_ID={3},Rank="{4}";'
                             .format(table, root, 'root', ID, rank))
            print('\t|\t'.join([table, root, 'root', ID, rank]))
            for OTU, parent in self.walk(startNode, parent=root):
                exists, IDs = self.check_exist(OTU, parent, table)
                if exists:
                    print("Name '{0}' and parent_name '{1}' aren't unique in {2} \\ {3}".format(OTU, parent, root, table))
                    exIDs, new = self._get_non_unique_tax_Ids(root, OTU, parent)
                    allIDs = []
                    for ID in exIDs + IDs:
                        if ID not in allIDs:
                            allIDs.append(str(ID))
                    IDs = ' '.join(allIDs)
                    if new:
                        cmd = 'INSERT INTO NonUniques SET Root="{0}",' +\
                              'Name="{1}",ParentName="{2}",tax_IDs="{3}";'\
                              .format(root, OTU, parent, IDs)
                        self.cur.execute(cmd)
                    else:
                        cmd = 'UPDATE NonUniques SET tax_IDs="{0}" WHERE ' +\
                              'Name="{1}" AND ParentName="{2}" AND Root="{3}";'\
                              .format(allIDs, OTU, parent, root)
                        self.cur.execute()
                    print('\t|\t'.join(['NonUniques', OTU, parent, ', '.join(allIDs)]))
                else:
                    ID, rank = self.get_Ids(OTU, parent)
                    cmd = 'INSERT INTO {0} SET Name="{1}",ParentName="{2}",' +\
                          'tax_ID={3},Rank="{4}";'\
                          .format(table, OTU, parent, ID, rank)
                    self.cur.execute(cmd)
                    print('\t|\t'.join([table, OTU, parent, ID, rank]))
            self.cnx.commit()

    def check_exist(self, OTU, parent, table):
        cmd = 'SELECT tax_ID FROM {0} WHERE Name="{1}" AND ParentName="{2}";'\
              .format(table, OTU, parent)
        self.cur.execute(cmd)
        results = self.rows_to_list(self.cur.fetchall())
        if len(results) == 0:
            return (False, None)
        else:
            return (True, results)

    def _write_text_files(self):
        for table in self.tables:
            with file(os.path.join(CONFIG.top, table+'.sql'), 'w') as outfile:
                print('Dumping {0} to {0}.sql'.format( table ))
                cmd = ['mysqldump', '-u', self.user, '-p'+self.passwd,
                       'taxonomies', table, '>', table+'.sql']
                subprocess.Popen(cmd, shell=True, stdout=outfile, stderr=subprocess.PIPE)

    def _clean_up(self):
        self._write_text_files()
        self.cnx.close()

    def automate(self):
        try:
            self._update_all()
        finally:
            self._clean_up()

if __name__ == '__main__':
    print(CONFIG.user)
    #print(CONFIG.passwd)
    if len(sys.argv) < 2:
        print("\nNot enough arguments given. What to do??")
        print("Either reload the NCBI taxonomy database with 'NCBI'")
        print("or compare all ARB definitions to EMBL's hosted at NCBI, with 'compare'")
        print("ie. Either: \n\t> python link_EMBL_taxonomys.py NCBI")
        print("or: \n\t> python link_EMBL_taxonomys.py compare\n")
        exit(1)
    elif sys.argv[1] == 'NCBI':
        DBCNX = NcbiTaxDb()
        DBCNX.automate()
    elif sys.argv[1] == 'compare':
        trans = Translate()
        trans.automate()
    elif sys.argv[1] == '--help':
        print(__doc__)
    else:
        raise KeyError("Wrong argument. Try 'NCBI', 'compare' or '--help'")
